/*
  Copyright (C) 2011, Kyungjoo Kim
  
  This file is part of LINAL (LINear ALgebra)
  
  All rights reserved.

  This library is free software; you can redistribute it and/or
  modify it under the terms of the GNU Lesser General Public
  License as published by the Free Software Foundation; either
  version 2.1 of the License, or (at your option) any later version.
  
  This library is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.
  
  You should have received a copy of the GNU Lesser General Public
  License along with this library; if not, write to the Free Software
  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  Also add information on how to contact you by electronic and paper mail.
  
  
  Kyungjoo Kim
  iamkyungjoo@gmail.com
*/
#include "linal/common.hxx"
#include "linal/util.hxx"

#include "linal/format.hxx"

#include "linal/operation.hxx"

#ifdef LINAL_DISPATCHER_ENABLE
#include "linal/device/operation.hxx"
#include "linal/device/device.hxx"
#include "linal/device/cache.hxx"
#include "linal/dispatcher.hxx"
#else
#include "linal/device/host/operation.hxx"
#endif

namespace linal {
  /*!
    General matrix matrix multiplication : Only update triangular part
  */
    
  int gemm_l( int transa, int transb,
              FLA_Obj alpha, Hier_ A, Hier_ B,
              FLA_Obj beta,  Hier_ C ) {

    LINAL_ERROR_DETAIL( check_trans_all( transa ),
                        "%s\n %s %d\n",
                        LINAL_ERROR_CHECK_TRANS,
                        "transa = ", transa );

    LINAL_ERROR_DETAIL( check_trans_all( transb ),
                        "%s\n %s %d\n",
                        LINAL_ERROR_CHECK_TRANS,
                        "transb = ", transb );
      
    void
      *buffer_A = A.get_buffer(),
      *buffer_B = B.get_buffer(),
      *buffer_C = C.get_buffer();
      
    int
      cs_A = A.get_cs(),
      cs_B = B.get_cs(),
      cs_C = C.get_cs();
      
    scal( beta, C );

    switch (transa) {
    case FLA_NO_TRANSPOSE:
    case FLA_CONJ_NO_TRANSPOSE:
      {
        switch (transb) {
        case FLA_NO_TRANSPOSE:
        case FLA_CONJ_NO_TRANSPOSE:
          {
            for (int p=0;p<A.get_n();++p) {
              for (int k2=0;k2<C.get_n();++k2) {
                for (int k1=k2;k1<C.get_m();++k1) {
                    
#pragma omp task firstprivate ( k1, k2, p )
                  {
#ifdef LINAL_DISPATCHER_ENABLE
                    get_dispatcher()->gemm( transa, transb,
                                            alpha,
                                            elem<FLA_Obj>(buffer_A, cs_A, 1, k1,p),
                                            elem<FLA_Obj>(buffer_B, cs_B, 1, p,k2),
                                            FLA_ONE,
                                            elem<FLA_Obj>(buffer_C, cs_C, 1, k1,k2) );

#else
                    host::internal::gemm( transa, transb,
                                          alpha,
                                          elem<FLA_Obj>(buffer_A, cs_A, 1, k1,p),
                                          elem<FLA_Obj>(buffer_B, cs_B, 1, p,k2),
                                          FLA_ONE,
                                          elem<FLA_Obj>(buffer_C, cs_C, 1, k1,k2) );
#endif
                  }
                }
              }
#pragma omp taskwait
            }
            break;
          }
        case FLA_TRANSPOSE:
        case FLA_CONJ_TRANSPOSE:
          {
            for (int p=0;p<A.get_n();++p) {
              for (int k2=0;k2<C.get_n();++k2) {
                for (int k1=k2;k1<C.get_m();++k1) {
                    
#pragma omp task firstprivate ( k1, k2, p )
                  {
#ifdef LINAL_DISPATCHER_ENABLE
                    get_dispatcher()->gemm( transa, transb,
                                            alpha,
                                            elem<FLA_Obj>(buffer_A, cs_A, 1, k1,p),
                                            elem<FLA_Obj>(buffer_B, cs_B, 1, k2,p),
                                            FLA_ONE,
                                            elem<FLA_Obj>(buffer_C, cs_C, 1, k1,k2) );

#else
                    host::internal::gemm( transa, transb,
                                          alpha,
                                          elem<FLA_Obj>(buffer_A, cs_A, 1, k1,p),
                                          elem<FLA_Obj>(buffer_B, cs_B, 1, k2,p),
                                          FLA_ONE,
                                          elem<FLA_Obj>(buffer_C, cs_C, 1, k1,k2) );
#endif
                  }
                }
              }
#pragma omp taskwait
            }
            break;
          }
        }
        break;
      }
    case FLA_TRANSPOSE:
    case FLA_CONJ_TRANSPOSE:
      {
        switch (transb) {
        case FLA_NO_TRANSPOSE:
        case FLA_CONJ_NO_TRANSPOSE:
          {
            for (int p=0;p<A.get_m();++p) {
              for (int k2=0;k2<C.get_n();++k2) {
                for (int k1=k2;k1<C.get_m();++k1) {
                    
#pragma omp task firstprivate ( k1, k2, p )
                  {
#ifdef LINAL_DISPATCHER_ENABLE
                    get_dispatcher()->gemm( transa, transb,
                                            alpha,
                                            elem<FLA_Obj>(buffer_A, cs_A, 1, p,k1),
                                            elem<FLA_Obj>(buffer_B, cs_B, 1, p,k2),
                                            FLA_ONE,
                                            elem<FLA_Obj>(buffer_C, cs_C, 1, k1,k2) );

#else
                    host::internal::gemm( transa, transb,
                                          alpha,
                                          elem<FLA_Obj>(buffer_A, cs_A, 1, p,k1),
                                          elem<FLA_Obj>(buffer_B, cs_B, 1, p,k2),
                                          FLA_ONE,
                                          elem<FLA_Obj>(buffer_C, cs_C, 1, k1,k2) );
#endif
                  }
                }
              }
#pragma omp taskwait
            }
            break;
          }
        case FLA_TRANSPOSE:
        case FLA_CONJ_TRANSPOSE:
          {
            for (int p=0;p<A.get_m();++p) {
              for (int k2=0;k2<C.get_n();++k2) {
                for (int k1=k2;k1<C.get_m();++k1) {
                    
#pragma omp task firstprivate ( k1, k2, p )
                  {
#ifdef LINAL_DISPATCHER_ENABLE
                    get_dispatcher()->gemm( transa, transb,
                                            alpha,
                                            elem<FLA_Obj>(buffer_A, cs_A, 1, p,k1),
                                            elem<FLA_Obj>(buffer_B, cs_B, 1, k2,p),
                                            FLA_ONE,
                                            elem<FLA_Obj>(buffer_C, cs_C, 1, k1,k2) );

#else
                    host::internal::gemm( transa, transb,
                                          alpha,
                                          elem<FLA_Obj>(buffer_A, cs_A, 1, p,k1),
                                          elem<FLA_Obj>(buffer_B, cs_B, 1, k2,p),
                                          FLA_ONE,
                                          elem<FLA_Obj>(buffer_C, cs_C, 1, k1,k2) );
#endif
                  }
                }
              }
#pragma omp taskwait
            }
            break;
          }
        }
        break;
      }
    default:
      {
        break;
      }
    }

    return true;
  }
    
  int gemm_u( int transa, int transb,
              FLA_Obj alpha, Hier_ A, Hier_ B,
              FLA_Obj beta,  Hier_ C ) {

    LINAL_ERROR_DETAIL( check_trans_all( transa ),
                        "%s\n %s %d\n",
                        LINAL_ERROR_CHECK_TRANS,
                        "transa = ", transa );

    LINAL_ERROR_DETAIL( check_trans_all( transb ),
                        "%s\n %s %d\n",
                        LINAL_ERROR_CHECK_TRANS,
                        "transb = ", transb );

    void
      *buffer_A = A.get_buffer(),
      *buffer_B = B.get_buffer(),
      *buffer_C = C.get_buffer();

    int
      cs_A = A.get_cs(),
      cs_B = B.get_cs(),
      cs_C = C.get_cs();

    scal( beta, C );

    switch (transa) {
    case FLA_NO_TRANSPOSE:
    case FLA_CONJ_NO_TRANSPOSE:
      {
        switch (transb) {
        case FLA_NO_TRANSPOSE:
        case FLA_CONJ_NO_TRANSPOSE:
          {
            for (int p=0;p<A.get_n();++p) {
              for (int k2=0;k2<C.get_n();++k2) {
                for (int k1=0;k1<(k2+1);++k1) {
                    
#pragma omp task firstprivate ( k1, k2, p )
                  {
#ifdef LINAL_DISPATCHER_ENABLE
                    get_dispatcher()->gemm( transa, transb,
                                            alpha,
                                            elem<FLA_Obj>(buffer_A, cs_A, 1, k1,p),
                                            elem<FLA_Obj>(buffer_B, cs_B, 1, p,k2),
                                            FLA_ONE,
                                            elem<FLA_Obj>(buffer_C, cs_C, 1, k1,k2) );

#else
                    host::internal::gemm( transa, transb,
                                          alpha,
                                          elem<FLA_Obj>(buffer_A, cs_A, 1, k1,p),
                                          elem<FLA_Obj>(buffer_B, cs_B, 1, p,k2),
                                          FLA_ONE,
                                          elem<FLA_Obj>(buffer_C, cs_C, 1, k1,k2) );
#endif
                  }
                }
              }
#pragma omp taskwait
            }
            break;
          }
        case FLA_TRANSPOSE:
        case FLA_CONJ_TRANSPOSE:
          {
            for (int p=0;p<A.get_n();++p) {
              for (int k2=0;k2<C.get_n();++k2) {
                for (int k1=0;k1<(k2+1);++k1) {
                    
#pragma omp task firstprivate ( k1, k2, p )
                  {
#ifdef LINAL_DISPATCHER_ENABLE
                    get_dispatcher()->gemm( transa, transb,
                                            alpha,
                                            elem<FLA_Obj>(buffer_A, cs_A, 1, k1,p),
                                            elem<FLA_Obj>(buffer_B, cs_B, 1, k2,p),
                                            FLA_ONE,
                                            elem<FLA_Obj>(buffer_C, cs_C, 1, k1,k2) );

#else
                    host::internal::gemm( transa, transb,
                                          alpha,
                                          elem<FLA_Obj>(buffer_A, cs_A, 1, k1,p),
                                          elem<FLA_Obj>(buffer_B, cs_B, 1, k2,p),
                                          FLA_ONE,
                                          elem<FLA_Obj>(buffer_C, cs_C, 1, k1,k2) );
#endif
                  }
                }
              }
#pragma omp taskwait
            }
            break;
          }
        }
        break;
      }
    case FLA_TRANSPOSE:
    case FLA_CONJ_TRANSPOSE:
      {
        switch (transb) {
        case FLA_NO_TRANSPOSE:
        case FLA_CONJ_NO_TRANSPOSE:
          {
            for (int p=0;p<A.get_m();++p) {
              for (int k2=0;k2<C.get_n();++k2) {
                for (int k1=0;k1<(k2+1);++k1) {
                    
#pragma omp task firstprivate ( k1, k2, p )
                  {
#ifdef LINAL_DISPATCHER_ENABLE
                    get_dispatcher()->gemm( transa, transb,
                                            alpha,
                                            elem<FLA_Obj>(buffer_A, cs_A, 1, p,k1),
                                            elem<FLA_Obj>(buffer_B, cs_B, 1, p,k2),
                                            FLA_ONE,
                                            elem<FLA_Obj>(buffer_C, cs_C, 1, k1,k2) );

#else
                    host::internal::gemm( transa, transb,
                                          alpha,
                                          elem<FLA_Obj>(buffer_A, cs_A, 1, p,k1),
                                          elem<FLA_Obj>(buffer_B, cs_B, 1, p,k2),
                                          FLA_ONE,
                                          elem<FLA_Obj>(buffer_C, cs_C, 1, k1,k2) );
#endif
                  }
                }
              }
#pragma omp taskwait
            }
            break;
          }
        case FLA_TRANSPOSE:
        case FLA_CONJ_TRANSPOSE:
          {
            for (int p=0;p<A.get_m();++p) {
              for (int k2=0;k2<C.get_n();++k2) {
                for (int k1=0;k1<(k2+1);++k1) {
                    
#pragma omp task firstprivate ( k1, k2, p )
                  {
#ifdef LINAL_DISPATCHER_ENABLE
                    get_dispatcher()->gemm( transa, transb,
                                            alpha,
                                            elem<FLA_Obj>(buffer_A, cs_A, 1, p,k1),
                                            elem<FLA_Obj>(buffer_B, cs_B, 1, k2,p),
                                            FLA_ONE,
                                            elem<FLA_Obj>(buffer_C, cs_C, 1, k1,k2) );

#else
                    host::internal::gemm( transa, transb,
                                          alpha,
                                          elem<FLA_Obj>(buffer_A, cs_A, 1, p,k1),
                                          elem<FLA_Obj>(buffer_B, cs_B, 1, k2,p),
                                          FLA_ONE,
                                          elem<FLA_Obj>(buffer_C, cs_C, 1, k1,k2) );
#endif
                  }
                }
              }
#pragma omp taskwait
            }
            break;
          }
        }
        break;
      }
    default:
      {
        break;
      }
    }

    return true;
  }
}
